import gradio as gr
import torch
import numpy as np
import json
import os
from PIL import Image
from torchvision import transforms
import torch.nn.functional as F
import platform

from model import get_model

# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载类别
def load_classes():
    """加载鸟类分类的类别"""
    dataset_path = "images/"
    classes = []
    
    if os.path.exists(dataset_path):
        # 从数据集目录加载类别
        class_dirs = sorted([d for d in os.listdir(dataset_path) if os.path.isdir(os.path.join(dataset_path, d))])
        classes = class_dirs
    elif os.path.exists('classes.txt'):
        # 从类别文件加载
        with open('classes.txt', 'r') as f:
            classes = [line.strip() for line in f.readlines()]
    else:
        # 如果找不到类别信息，使用默认类别
        print("警告: 未找到类别信息，使用默认类别名称")
        classes = [f"类别_{i}" for i in range(200)]
    
    return classes

# 格式化类别名称
def format_class_name(class_name):
    """格式化类名，移除前缀和替换连字符"""
    if '.' in class_name and class_name[0].isdigit():
        # 移除类似 "001." 的前缀
        parts = class_name.split('.')
        if len(parts) > 1:
            class_name = '.'.join(parts[1:])
    return class_name.replace('-', ' ').replace('_', ' ')

# 加载模型
def load_model(model_path='model/best_model.pth', num_classes=200):
    """加载预训练的鸟类分类模型"""
    try:
        model = get_model(num_classes)
        
        # 尝试加载模型权重
        if os.path.exists(model_path):
            checkpoint = torch.load(model_path, map_location=device)
            # 兼容两种保存格式
            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
                model.load_state_dict(checkpoint['model_state_dict'])
            else:
                model.load_state_dict(checkpoint)
            
            print(f"成功加载模型: {model_path}")
        else:
            print(f"错误: 模型文件 {model_path} 不存在")
            return None
        
        model = model.to(device)
        model.eval()
        return model
    except Exception as e:
        print(f"加载模型时出错: {str(e)}")
        return None

# 图像预处理
def preprocess_image(image):
    """预处理图像用于模型输入"""
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    # 确保图像是PIL格式
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # 应用变换
    image_tensor = transform(image).unsqueeze(0).to(device)
    return image_tensor

# 预测函数
def predict_image(image, top_k=5):
    """使用模型预测图像类别"""
    if image is None:
        return {"错误: 未提供图像": 1.0}
    
    # 加载类别和模型
    classes = load_classes()
    model = load_model(num_classes=len(classes))
    
    if model is None:
        return {"错误: 无法加载模型": 1.0}
    
    # 预处理图像
    image_tensor = preprocess_image(image)
    
    # 模型推理
    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = F.softmax(outputs, dim=1)
    
    # 获取前K个预测结果
    top_probs, top_indices = torch.topk(probabilities, top_k)
    
    # 转换为numpy数组
    top_probs = top_probs.squeeze().cpu().numpy()
    top_indices = top_indices.squeeze().cpu().numpy()
    
    # 创建结果字典
    results = {}
    for i in range(top_k):
        class_name = format_class_name(classes[top_indices[i]])
        results[class_name] = float(top_probs[i])
    
    return results

# 获取系统信息
def get_system_info():
    """获取简要系统信息"""
    info = f"""
    - 操作系统: {platform.system()} {platform.release()}
    - Python版本: {platform.python_version()}
    - PyTorch版本: {torch.__version__}
    - 设备: {"GPU" if torch.cuda.is_available() else "CPU"}
    """
    if torch.cuda.is_available():
        info += f"- GPU: {torch.cuda.get_device_name(0)}"
    
    return info

# 创建Gradio界面
def main():
    # 加载类别和模型信息
    classes = load_classes()
    model_path = 'model/best_model.pth'
    model_exists = os.path.exists(model_path)
    
    # 创建标题和描述
    title = "🦜 鸟类识别系统"
    
    description = f"""
    ## 基于ResNet101的鸟类识别系统
    
    该系统可以识别200种不同的鸟类。上传一张鸟类图片，系统将给出最可能的鸟类种类及其置信度。
    
    ### 系统信息
    {get_system_info()}
    
    ### 模型状态
    - 模型路径: {model_path}
    - 模型状态: {"已加载 ✓" if model_exists else "未找到 ✗"}
    - 支持类别数: {len(classes)}
    """
    
    # 创建示例列表
    examples = []
    example_dir = os.path.join(os.getcwd(), 'example_images')
    if os.path.exists(example_dir):
        example_files = [f for f in os.listdir(example_dir) 
                        if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        examples = [os.path.join(example_dir, f) for f in example_files[:5]]  # 最多5个示例
    
    # 创建界面
    iface = gr.Interface(
        fn=predict_image,
        inputs=gr.Image(type="pil", label="上传鸟类图片"),
        outputs=gr.Label(label="识别结果", num_top_classes=5),
        title=title,
        description=description,
        examples=examples,
        theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue),
        allow_flagging="never"
    )
    
    # 启动服务器
    iface.launch(share=False)

if __name__ == "__main__":
    main()
